import asyncio
import re

from agents.base_agent import BaseAgent
import json
from typing import List, Dict
from copy import deepcopy

from my_utils.format_parser import extract_simple_name


class RevisionManager:
    """修订管理类，负责控制整个修订流程"""

    def __init__(self, score_threshold: int = 95, min_improvement: int = 5):
        self.score_threshold = score_threshold  # 评分阈值（高于此分数无需修改）
        self.min_improvement = min_improvement  # 最小改进分数（低于此分数视为修改效果不明显）

    def check_need_further_revision(self, old_evaluations: List[Dict],
                                    new_evaluations: List[Dict]) -> bool:
        """ 检查是否需要继续修改
        Args:
            old_evaluations: 上一轮评价结果
            new_evaluations: 新的评价结果 """
        if not old_evaluations:  # 第一轮修改
            return any(e['score'] < self.score_threshold for e in new_evaluations)

        # 检查是否有显著改进
        improvements = []
        for old_eval, new_eval in zip(old_evaluations, new_evaluations):
            improvement = new_eval['score'] - old_eval['score']
            improvements.append(improvement)

        # 如果任何章节仍然分数过低且有明显改进，继续修改
        for score, improvement in zip([e['score'] for e in new_evaluations], improvements):
            # if score < self.score_threshold and improvement >= self.min_improvement:
            #     return True
            if improvement >= self.min_improvement:
                return True

        return False

    def merge_revisions(self, doc: Dict, revised_chapters: List[Dict]) -> Dict:
        """ 合并修改后的章节 """
        doc_copy = deepcopy(doc)

        # 更新章节内容
        for i, chapter in enumerate(revised_chapters):
            if i < len(doc_copy['children']):
                doc_copy['children'][i]['content'] = chapter['content']
                # 合并新增的引用
                if 'new_refs' in chapter and chapter['new_refs']:
                    # 获取原有的引用
                    # existing_refs = set(doc_copy['children'][i]['refs'])
                    new_refs = [paper_chunk_id.split('-')[0] for paper_chunk_id in chapter['new_refs']]
                    # doc_copy['children'][i]['refs'] = list(existing_refs.union(new_refs))
                    doc_copy['children'][i]['refs'] = new_refs
        return doc_copy


class DocumentReviser():
    """文档修订主类，协调整个修订过程"""

    def __init__(self, max_iterations: int = 1):
        """ 初始化文档修订器 """
        self.max_iterations = max_iterations  # 最大修改轮次
        self.evaluator = None
        self.reviser = None
        self.manager = None
        self.ref_adder = None

    def setup_agents(self):
        """初始化所有需要的组件"""
        self.manager = RevisionManager()
        self.ref_adder = BaseAgent()

    def revise_document(self, doc: Dict) -> Dict:
        """ 文档修订主流程 """
        if not all([self.evaluator, self.reviser, self.manager, self.ref_adder]):
            self.setup_agents()

        current_doc = deepcopy(doc)
        current_chapters = current_doc['children']
        previous_evaluations = []

        for iteration in range(self.max_iterations):
            print(f"开始第 {iteration + 1} 轮修改...")

            # 评价当前版本
            current_evaluations = self.evaluator.evaluate_chapters(current_chapters)

            # 修改章节
            revised_chapters = self.reviser.revise_chapters(
                current_doc, current_evaluations)

            # 更新文档
            current_doc = self.manager.merge_revisions(current_doc, revised_chapters)
            current_chapters = current_doc['children']

            # 保存中间结果
            with open(f"examples_data/revision_iteration_{iteration + 1}.json", "w", encoding='utf-8') as f:
                json.dump(current_doc, f, ensure_ascii=False, indent=4)

            previous_evaluations = current_evaluations

            print(f"第 {iteration + 1} 轮修改完成。")

        # 将current_doc['children']列表中的每个字典的content属性，拼接起来
        current_doc["content"] = "\n\n".join([chapter['content'] for chapter in current_doc['children']])

        return current_doc

    @staticmethod
    def revise_paper_references(text):
        """ 
        修改文本中的论文引用格式:
        1. 将(paper_XXX)格式转换为<sup>paper_XXX</sup>
        2. 将(<sup>paper_XXX</sup>)格式转换为<sup>paper_XXX</sup>
        """
        # 第一步：处理(paper_XXX)格式
        pattern1 = r'\(paper_(.*?)\)'
        def replace_with_sup(match):
            paper_id = match.group(1)
            return f"<sup>paper_{paper_id}</sup>"
        
        result = re.sub(pattern1, replace_with_sup, text)
        
        # 第二步：处理(<sup>paper_XXX</sup>)格式
        pattern2 = r'\(<sup>(paper_.*?)</sup>\)'
        result = re.sub(pattern2, r'<sup>\1</sup>', result)
        
        return result

    @staticmethod
    def wrap_paper_ids(text):
        """ 提取文本中的 paper_id，并确保它们被 <sup></sup> 标签包裹。 """

        # 正则表达式：匹配 paper_id (paper_后面跟着一个或多个数字)
        paper_id_pattern = r'paper_\d+'

        def replace_paper_id(match):
            paper_id = match.group(0)  # 匹配到的 paper_id

            # 检查是否已经被 <sup></sup> 包裹
            if not (text.rfind('<sup>', 0, match.start()) > text.rfind('</sup>', 0, match.start())):
                # 使用rfind 从paper_id往前找，找到<sup>的索引大于</sup>的索引
                return f'this paper <sup>{paper_id}</sup>'
            else:
                return paper_id

        # 使用 re.sub() 进行替换
        modified_text = re.sub(paper_id_pattern, replace_paper_id, text)
        return modified_text

    @staticmethod
    def extract_and_replace_paper_ids(text):
        def remove_non_alphanumeric(input_string):
            # 使用正则表达式匹配所有非字母和非数字的字符，并替换为空字符串
            return re.sub(r'[^0-9]', '', input_string)

        # 使用正则表达式提取所有<sup>paper_id</sup>中的paper_id
        paper_ids = re.findall(r'<sup>(.*?)</sup>', text)
        # 提取出包含'(paper_(.*?))'的字符串'(paper_(.*?))'，然后将括号转为'<sup></sup>'



        paper_ids += re.findall(r'\(paper_(.*?)\)', text)

        # 去重并保持顺序
        unique_paper_ids = sorted(set(paper_ids), key=paper_ids.index)
        unique_paper_ids = [remove_non_alphanumeric(paper_id) for paper_id in unique_paper_ids]
        # 创建一个映射，将每个unique_paper_id映射到一个序号
        paper_id_to_number = {paper_id: str(i + 1) for i, paper_id in enumerate(unique_paper_ids)}

        # 替换文本中的paper_id为相应的序号
        def replace_paper_id(match):
            raw = match.group(1)
            raw = remove_non_alphanumeric(raw)
            return f'<sup>{paper_id_to_number[raw]}</sup>'

        # 使用正则替换文本中的paper_id
        result_text = re.sub(r'<sup>(.*?)</sup>', replace_paper_id, text)

        return result_text, paper_id_to_number


def revise_paper(filled_outline_recursive_dict_copy: Dict, ref_id_list: List, selected_chunks: Dict) -> str:
    """ 处理并格式化论文内容，包括引用编号的替换和参考文献列表的生成 """
    # 初始化文档修订器
    reviser = DocumentReviser()
    revise_document = filled_outline_recursive_dict_copy

    # 将文本中的引用标识(<sup>paper_id</sup>)转换为数字编号
    revise_document['content'] = reviser.revise_paper_references(revise_document['content'])
    revise_document['content'] = reviser.wrap_paper_ids(revise_document['content'])
    result_text, paper_id_to_number = reviser.extract_and_replace_paper_ids(revise_document['content'])

    # 提取实际使用的引用编号列表
    ref_id_used_list = list(paper_id_to_number.keys())

    # 将引用编号转换为对应的索引值（减1是因为ref_id从1开始）
    ref_indices = [int(ref_id) - 1 for ref_id in ref_id_used_list if ref_id.strip()]

    # 根据索引获取实际的文献ID
    unique_paper_chunk_ids = []
    for ref_index in ref_indices:
        if ref_index < len(ref_id_list):
            # 如果索引有效，直接使用对应的文献ID
            unique_paper_chunk_ids.append(ref_id_list[ref_index])
        else:
            # 索引无效时的备选策略：使用相邻的文献ID
            if unique_paper_chunk_ids:
                # 已有引用时，选择相邻的文献
                prev_chunk = unique_paper_chunk_ids[-1]
                prev_index = ref_id_list.index(prev_chunk)
                next_index = prev_index + 1 if prev_index < len(ref_id_list) - 1 else prev_index - 1
                unique_paper_chunk_ids.append(ref_id_list[next_index])
            else:
                # 没有任何引用时，使用第一个文献
                unique_paper_chunk_ids.append(ref_id_list[0])

    # 提取文献的主ID（去除chunk标识）
    paper_ids = [chunk_id.split('-')[0] for chunk_id in unique_paper_chunk_ids]

    # 收集被引用文献的详细信息
    references_info = []
    for paper_id in paper_ids:
        paper_entity = selected_chunks[paper_id]['entity']
        references_info.append({
            "paper_title": paper_entity['paper_title'],
            "chunk_id": paper_entity['chunk_id'],
            "paper_database": paper_entity['original_filename'],
            "year": paper_entity['year']
        })

    # 生成参考文献列表
    references_section = "# References\n"
    for idx, ref in enumerate(references_info, 1):
        venue = ref['paper_database'].split('.')[0]  # 提取会议/期刊名称
        simple_name = extract_simple_name(venue)
        references_section += (f"[{idx}] {ref['paper_title']}, {simple_name}, {ref['year']}, "
                               f"chunk {ref['chunk_id']}\n\n")

    # 组装最终文档（标题 + 正文 + 参考文献）
    final_document = (f"{revise_document['title']}\n"
                      f"{result_text}\n\n"
                      f"{references_section}")

    return final_document

